package io.kiki.sba.registry.metrics;


import com.codahale.metrics.Gauge;
import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricRegistry;
import io.kiki.sba.registry.task.KeyedThreadPoolExecutor;
import io.kiki.sba.registry.util.ConcurrentUtils;
import io.kiki.sba.registry.util.WakeUpLoopExecuteTask;
import io.kiki.stack.netty.protocol.ProtocolImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.regex.Pattern;


public class TaskMetrics {
    private static final Logger logger = LoggerFactory.getLogger(TaskMetrics.class);
    private static final TaskMetrics instance = new TaskMetrics();
    private static final String EXECUTE_NAME_SEP = Pattern.quote(".");
    private static boolean boltRegistered = false;
    private final MetricRegistry metrics = new MetricRegistry();
    private final Set<String> executors = Collections.newSetFromMap(new ConcurrentHashMap<>());

    private TaskMetrics() {
        ExecutorMetricsWatchDog executorMetricsWatchDog = new ExecutorMetricsWatchDog();
        ConcurrentUtils.createDaemonThread("executorMetrics", executorMetricsWatchDog).start();
    }

    public static TaskMetrics getInstance() {
        return instance;
    }


    public static void setBoltRegistered(boolean boltRegistered) {
        TaskMetrics.boltRegistered = boltRegistered;
    }

    public MetricRegistry getMetricRegistry() {
        return this.metrics;
    }

    public synchronized void registerBolt() {
        if (boltRegistered) {
            return;
        }
        ThreadPoolExecutor boltDefaultExecutor = (ThreadPoolExecutor) ProtocolImpl.protocol.getCommandHandler().getDefaultExecutor();
        registerThreadExecutor("BoltDefaultExecutor", boltDefaultExecutor);
        boltRegistered = true;
    }

    public void registerKeyThreadExecutor(String executorName, KeyedThreadPoolExecutor executor) {
        registerIfAbsent(MetricRegistry.name(executorName, "queue"), (Gauge<Integer>) executor::getQueueSize);

        registerIfAbsent(MetricRegistry.name(executorName, "poolSize"), (Gauge<Integer>) executor::getCoreSize);

        registerIfAbsent(MetricRegistry.name(executorName, "active"), (Gauge<Integer>) executor::getActiveCount);

        registerIfAbsent(MetricRegistry.name(executorName, "completed"), (Gauge<Long>) executor::getCompletedTaskCount);

        registerIfAbsent(MetricRegistry.name(executorName, "task"), (Gauge<Long>) executor::getTaskCount);
        executors.add(executorName);
    }

    public void registerThreadExecutor(String executorName, ThreadPoolExecutor executor) {
        registerIfAbsent(MetricRegistry.name(executorName, "queue"), (Gauge<Integer>) () -> executor.getQueue().size());

        registerIfAbsent(MetricRegistry.name(executorName, "poolSize"), (Gauge<Integer>) executor::getPoolSize);

        registerIfAbsent(MetricRegistry.name(executorName, "active"), (Gauge<Integer>) executor::getActiveCount);

        registerIfAbsent(MetricRegistry.name(executorName, "completed"), (Gauge<Long>) executor::getCompletedTaskCount);

        registerIfAbsent(MetricRegistry.name(executorName, "task"), (Gauge<Long>) executor::getTaskCount);
        executors.add(executorName);
    }

    private boolean registerIfAbsent(String metricName, Metric metric) {
        Map<String, Metric> metricMap = metrics.getMetrics();
        if (metricMap.containsKey(metricName)) {
            logger.warn("executor.metric exists {}", metricName);
            return false;
        }
        metrics.register(metricName, metric);
        return true;
    }

    class ExecutorMetricsWatchDog extends WakeUpLoopExecuteTask {

        @Override
        public void _execute_() {
            for (String executorName : executors) {
                Map<String, Gauge> map = metrics.getGauges((name, value) -> name.contains(".") && name.split(EXECUTE_NAME_SEP)[0].equals(executorName));
                StringBuilder sb = new StringBuilder();
                sb.append(executorName);
                map.forEach((key, gauge) -> {
                    String name = key.substring(executorName.length() + 1);
                    sb.append(", ").append(name).append(":").append(gauge.getValue());
                });
                logger.info(sb.toString());
            }
        }

        @Override
        public int getWaitingMillis() {
            return 5000;
        }
    }
}
